from typing import List

from sqlalchemy.orm import InstrumentedAttribute
from sqlalchemy.sql.elements import BinaryExpression

from sqlalchemy_tool.transcation_decorator import Transcation
from sqlalchemy_tool.model_utils import to_dict


class Manager:
    # 事务管理对象
    transcation = Transcation()
    model_class = None

    def __init__(self, model_obj=None, model_class=None):
        if model_obj:
            self.__update_model_obj__(model_obj)
        if model_class:
            self.__update_model_class__(model_class)

    def __update_model_class__(self, model_class):
        self.model_class = model_class
        self.table_name = model_class.__tablename__
        column_dict = dict()
        for k, v in model_class.__dict__.items():
            if isinstance(v, InstrumentedAttribute):
                column_dict[k] = v
        self.column_dict = column_dict
        self.columns = tuple([key for key in self.column_dict.keys()])

    def __update_model_obj__(self, model_obj):
        self.model_obj = model_obj
        self.model_class = type(self.model_obj)
        self.table_name = self.model_obj.__tablename__
        self.id = self.model_obj.id
        self.columns = tuple([key for key in self.model_obj.__column_dict__.keys()])

    @classmethod
    def get_session(cls):
        return cls.transcation.get_session()

    @transcation
    def add(self) -> bool:
        """插入一条记录"""
        session = self.get_session()
        session.add(self.model_obj)
        session.flush()
        return True

    @transcation
    def add_batch(self, model_list: list) -> bool:
        """插入（批量）"""
        session = self.get_session()
        session.bulk_save_objects(model_list)
        session.flush()
        return True

    @transcation
    def save(self, *criterion) -> bool:
        """
        存在更新记录，否则插入一条记录
        :param criterion: 更新条件如：CategoryInfo.id <= 280， 默认按照id更新
        :return:
        """
        session = self.get_session()
        if criterion:
            old_data_count = session.query(self.model_class).where(*criterion).count()
        else:
            old_data_count = session.query(self.model_class).where(self.model_class.id == self.model_obj.id).count()
        if old_data_count > 0:
            return self.update(*criterion)
        else:
            return self.add()

    @transcation
    def save_batch(self, model_list) -> bool:
        """
        批量存在更新记录，否则插入记录（通过id判断）
        :param model_list: Model数据集合
        :return:
        """
        session = self.get_session()
        add_model_list = []
        update_model_list = []
        for tmp_model in model_list:
            if session.query(type(tmp_model)).filter(type(tmp_model).id == tmp_model.id).count() > 0:
                update_model_list.append(tmp_model)
            else:
                add_model_list.append(tmp_model)
        if add_model_list:
            self.add_batch(add_model_list)
        if update_model_list:
            self.update_batch(update_model_list)
        return True

    @transcation
    def save_batch_by_where(self, model_list, criterion_fun=None) -> bool:
        """
        批量存在更新记录，否则插入记录
        :param model_list: Model数据集合
        :param criterion_fun: 通过接受model，返回判断条件如：type(tmp_model).id <= 280， 默认按照id更新
        :return:
        """
        session = self.get_session()
        for tmp_model in model_list:
            criterion = criterion_fun(tmp_model)
            if not criterion and tmp_model.id:
                criterion = type(tmp_model).id == tmp_model.id
            if type(criterion) == BinaryExpression:
                criterion = (criterion,)
            if criterion and session.query(type(tmp_model)).filter(*criterion).count() > 0:
                session.query(type(tmp_model)).filter(*criterion).update(to_dict(tmp_model))
            else:
                session.add(tmp_model)
        return True

    @transcation
    def update(self, *criterion) -> bool:
        """
        通过条件更新，默认通过id更新
        :param criterion: 更新条件如：CategoryInfo.id <= 280， 默认按照id更新
        :return: 是否更新成功
        """
        session = self.get_session()
        if criterion:
            session.query(self.model_class).filter(*criterion).update(
                to_dict(self.model_obj))
        else:
            session.query(self.model_class).filter(self.model_class.id == self.model_obj.id).update(
                to_dict(self.model_obj))
        return True

    @transcation
    def update_batch(self, model_list) -> bool:
        """批量通过id更新"""
        session = self.get_session()
        for tmp_model in model_list:
            session.query(type(tmp_model)).filter(type(tmp_model).id == self.model_obj.id).update(
                to_dict(tmp_model))
        return True

    def get_one(self, *criterion):
        """根据 where条件查询一条记录"""
        session = self.get_session()
        if criterion:
            return session.query(self.model_class).filter(*criterion).first()
        else:
            return session.query(self.model_class).first()

    def get_dict(self, *criterion):
        """根据 where条件查询一条记录，转为字典"""
        return to_dict(self.get_one(*criterion))

    def get_by_id(self, id_val: int):
        """根据 ID 查询"""
        session = self.get_session()
        return session.query(self.model_class).filter(self.model_class.id == id_val).first()

    def list(self, *criterion) -> List:
        """通过条件查询，默认查询所有"""
        session = self.get_session()
        if criterion:
            return session.query(self.model_class).filter(*criterion).all()
        else:
            return session.query(self.model_class).all()

    def list_by_ids(self, id_list: List[int]) -> List:
        """通过id集合查询"""
        return self.list(self.model_class.id.in_(id_list))

    def page(self, page_num, page_size, *criterion):
        """分页查询"""
        session = self.get_session()
        start_index = (page_num - 1) * page_size
        if criterion:
            return session.query(self.model_class).filter(*criterion).offset(start_index).limit(page_size).all()
        else:
            return session.query(self.model_class).offset(start_index).limit(page_size).all()

    def count(self, *criterion) -> int:
        """查询记录数"""
        session = self.get_session()
        if criterion:
            return session.query(self.model_class).filter(*criterion).count()
        else:
            return session.query(self.model_class).count()

    @transcation
    def delete(self, *criterion) -> bool:
        """
        删除记录
        :param criterion: 删除条件如：CategoryInfo.id <= 280， 默认按照id删除
        :return:
        """
        session = self.get_session()
        if criterion:
            session.query(self.model_class).filter(*criterion).delete()
        else:
            session.query(self.model_class).filter(self.model_class.id == self.model_obj.id).delete()
        return True

    @transcation
    def delete_by_ids(self, ids: List[int]) -> bool:
        """根据id批量删除"""
        session = self.get_session()
        session.query(self.model_class).filter(self.model_class.id.in_(ids)).delete()
        return True
